#!/bin/bash

# 5GC Testbed Health Check Script
# This script performs comprehensive health checks on all 5GC components

set -e  # Exit on any error

# Configuration
VM1="192.168.1.206"  # Core Network
VM2="192.168.1.219"  # UPF
VM3="192.168.1.221"  # gNodeB
VM4="192.168.1.205"  # UE

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Health check results
declare -A health_status
declare -A health_messages

# Logging function
log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] $1${NC}"
}

warn() {
    echo -e "${YELLOW}[$(date +'%Y-%m-%d %H:%M:%S')] WARNING: $1${NC}"
}

error() {
    echo -e "${RED}[$(date +'%Y-%m-%d %H:%M:%S')] ERROR: $1${NC}"
}

info() {
    echo -e "${BLUE}[$(date +'%Y-%m-%d %H:%M:%S')] INFO: $1${NC}"
}

# Function to check if a service is running
check_service() {
    local service=$1
    local vm=$2
    
    if ssh -o ConnectTimeout=5 $vm "systemctl is-active --quiet $service" 2>/dev/null; then
        return 0
    else
        return 1
    fi
}

# Function to check if a process is running
check_process() {
    local process=$1
    local vm=$2
    
    if ssh -o ConnectTimeout=5 $vm "pgrep -f '$process'" >/dev/null 2>&1; then
        return 0
    else
        return 1
    fi
}

# Function to check network connectivity
check_connectivity() {
    local vm=$1
    
    if ping -c 1 -W 5 $vm >/dev/null 2>&1; then
        return 0
    else
        return 1
    fi
}

# Function to check port connectivity
check_port() {
    local vm=$1
    local port=$2
    local protocol=${3:-tcp}
    
    if [ "$protocol" = "udp" ]; then
        if nc -zu $vm $port 2>/dev/null; then
            return 0
        else
            return 1
        fi
    else
        if nc -z $vm $port 2>/dev/null; then
            return 0
        else
            return 1
        fi
    fi
}

# Function to check SCTP port (for NGAP)
check_sctp_port() {
    local vm=$1
    local port=$2
    
    # Use netstat to check SCTP endpoints
    if ssh -o ConnectTimeout=5 $vm "netstat -ln | grep sctp | grep :$port" >/dev/null 2>&1; then
        return 0
    else
        return 1
    fi
}

# Function to record health status
record_status() {
    local component=$1
    local status=$2
    local message=$3
    
    health_status[$component]=$status
    health_messages[$component]=$message
}

# Print health check header
print_header() {
    echo
    echo "=========================================="
    echo "        5GC Testbed Health Check"
    echo "=========================================="
    echo "Date: $(date)"
    echo "=========================================="
    echo
}

# Print health check results
print_results() {
    echo
    echo "=========================================="
    echo "           Health Check Results"
    echo "=========================================="
    echo
    
    local total_checks=0
    local passed_checks=0
    local failed_checks=0
    local warning_checks=0
    
    for component in "${!health_status[@]}"; do
        total_checks=$((total_checks + 1))
        
        case ${health_status[$component]} in
            "PASS")
                echo -e "✓ ${GREEN}$component${NC}: ${health_messages[$component]}"
                passed_checks=$((passed_checks + 1))
                ;;
            "FAIL")
                echo -e "✗ ${RED}$component${NC}: ${health_messages[$component]}"
                failed_checks=$((failed_checks + 1))
                ;;
            "WARN")
                echo -e "⚠ ${YELLOW}$component${NC}: ${health_messages[$component]}"
                warning_checks=$((warning_checks + 1))
                ;;
        esac
    done
    
    echo
    echo "=========================================="
    echo "Summary:"
    echo "  Total Checks: $total_checks"
    echo -e "  ${GREEN}Passed: $passed_checks${NC}"
    echo -e "  ${YELLOW}Warnings: $warning_checks${NC}"
    echo -e "  ${RED}Failed: $failed_checks${NC}"
    echo "=========================================="
    
    if [ $failed_checks -gt 0 ]; then
        exit 1
    fi
}

# Check 1: VM Connectivity
check_vm_connectivity() {
    log "Checking VM connectivity..."
    
    for vm in $VM1 $VM2 $VM3 $VM4; do
        if check_connectivity $vm; then
            record_status "VM_$vm" "PASS" "VM is reachable"
        else
            record_status "VM_$vm" "FAIL" "VM is not reachable"
        fi
    done
}

# Check 2: Open5GS Services
check_open5gs_services() {
    log "Checking Open5GS services..."
    
    # Core services on VM1
    services=("open5gs-amfd" "open5gs-smfd" "open5gs-udmd" "open5gs-udrd" "open5gs-pcfd")
    
    for service in "${services[@]}"; do
        if check_service $service $VM1; then
            record_status "$service" "PASS" "Service is running"
        else
            record_status "$service" "FAIL" "Service is not running"
        fi
    done
    
    # UPF service on VM2
    if check_service "open5gs-upfd" $VM2; then
        record_status "open5gs-upfd" "PASS" "UPF service is running"
    else
        record_status "open5gs-upfd" "FAIL" "UPF service is not running"
    fi
}

# Check 3: UERANSIM Processes
check_ueransim_processes() {
    log "Checking UERANSIM processes..."
    
    if check_process "nr-gnb" $VM3; then
        record_status "gNodeB" "PASS" "gNodeB process is running"
    else
        record_status "gNodeB" "FAIL" "gNodeB process is not running"
    fi
    
    if check_process "nr-ue" $VM4; then
        record_status "UE" "PASS" "UE process is running"
    else
        record_status "UE" "FAIL" "UE process is not running"
    fi
}

# Check 4: Network Ports
check_network_ports() {
    log "Checking network ports..."
    
    # Core network ports (VM1)
    if check_port $VM1 7777; then
        record_status "Core_SBI" "PASS" "SBI port 7777 is listening"
    else
        record_status "Core_SBI" "FAIL" "SBI port 7777 is not listening"
    fi
    
    if check_port $VM1 38412; then
        record_status "Core_NGAP" "PASS" "NGAP port 38412 is listening"
    else
        record_status "Core_NGAP" "FAIL" "NGAP port 38412 is not listening"
    fi
    
    if check_port $VM1 8805; then
        record_status "Core_PFCP" "PASS" "PFCP port 8805 is listening"
    else
        record_status "Core_PFCP" "FAIL" "PFCP port 8805 is not listening"
    fi
    
    # UPF ports (VM2)
    if check_port $VM2 8805; then
        record_status "UPF_PFCP" "PASS" "UPF PFCP port 8805 is listening"
    else
        record_status "UPF_PFCP" "FAIL" "UPF PFCP port 8805 is not listening"
    fi
    
    if check_port $VM2 2152 udp; then
        record_status "UPF_GTPU" "PASS" "UPF GTP-U port 2152 is listening"
    else
        record_status "UPF_GTPU" "FAIL" "UPF GTP-U port 2152 is not listening"
    fi
}

# Check 5: SCTP Associations
check_sctp_associations() {
    log "Checking SCTP associations..."
    
    # Check gNodeB-AMF association
    if ssh $VM3 "ss -tan | grep '192.168.1.10:38412' | grep ESTAB" >/dev/null 2>&1; then
        record_status "gNB_AMF_SCTP" "PASS" "gNodeB-AMF SCTP association established"
    else
        record_status "gNB_AMF_SCTP" "FAIL" "gNodeB-AMF SCTP association not established"
    fi
    
    # Check SMF-UPF association
    if ssh $VM1 "ss -tan | grep '192.168.1.11:8805' | grep ESTAB" >/dev/null 2>&1; then
        record_status "SMF_UPF_SCTP" "PASS" "SMF-UPF SCTP association established"
    else
        record_status "SMF_UPF_SCTP" "FAIL" "SMF-UPF SCTP association not established"
    fi
}

# Check 6: Database Connectivity
check_database_connectivity() {
    log "Checking database connectivity..."
    
    if ssh $VM1 "mongo --eval 'db.runCommand(\"ping\")' --quiet" >/dev/null 2>&1; then
        record_status "MongoDB" "PASS" "MongoDB is accessible"
    else
        record_status "MongoDB" "FAIL" "MongoDB is not accessible"
    fi
}

# Check 7: System Resources
check_system_resources() {
    log "Checking system resources..."
    
    # Check CPU usage
    for vm in $VM1 $VM2 $VM3 $VM4; do
        cpu_usage=$(ssh $vm "top -bn1 | grep 'Cpu(s)' | awk '{print \$2}' | cut -d'%' -f1" 2>/dev/null || echo "0")
        if (( $(echo "$cpu_usage < 80" | bc -l) )); then
            record_status "CPU_$vm" "PASS" "CPU usage: ${cpu_usage}%"
        elif (( $(echo "$cpu_usage < 90" | bc -l) )); then
            record_status "CPU_$vm" "WARN" "High CPU usage: ${cpu_usage}%"
        else
            record_status "CPU_$vm" "FAIL" "Very high CPU usage: ${cpu_usage}%"
        fi
    done
    
    # Check memory usage
    for vm in $VM1 $VM2 $VM3 $VM4; do
        mem_usage=$(ssh $vm "free | grep Mem | awk '{printf \"%.1f\", \$3/\$2 * 100.0}'" 2>/dev/null || echo "0")
        if (( $(echo "$mem_usage < 80" | bc -l) )); then
            record_status "Memory_$vm" "PASS" "Memory usage: ${mem_usage}%"
        elif (( $(echo "$mem_usage < 90" | bc -l) )); then
            record_status "Memory_$vm" "WARN" "High memory usage: ${mem_usage}%"
        else
            record_status "Memory_$vm" "FAIL" "Very high memory usage: ${mem_usage}%"
        fi
    done
}

# Check 8: Recent Errors in Logs
check_recent_errors() {
    log "Checking for recent errors in logs..."
    
    # Check Open5GS logs for errors in the last hour
    error_count=$(ssh $VM1 "journalctl -u open5gs-* --since '1 hour ago' | grep -i error | wc -l" 2>/dev/null || echo "0")
    if [ $error_count -eq 0 ]; then
        record_status "Core_Logs" "PASS" "No recent errors in core logs"
    elif [ $error_count -lt 10 ]; then
        record_status "Core_Logs" "WARN" "$error_count recent errors in core logs"
    else
        record_status "Core_Logs" "FAIL" "$error_count recent errors in core logs"
    fi
    
    # Check UERANSIM logs for errors
    gnb_errors=$(ssh $VM3 "grep -i error gnb.log | tail -10 | wc -l" 2>/dev/null || echo "0")
    if [ $gnb_errors -eq 0 ]; then
        record_status "gNodeB_Logs" "PASS" "No recent errors in gNodeB logs"
    else
        record_status "gNodeB_Logs" "WARN" "$gnb_errors recent errors in gNodeB logs"
    fi
    
    ue_errors=$(ssh $VM4 "grep -i error ue.log | tail -10 | wc -l" 2>/dev/null || echo "0")
    if [ $ue_errors -eq 0 ]; then
        record_status "UE_Logs" "PASS" "No recent errors in UE logs"
    else
        record_status "UE_Logs" "WARN" "$ue_errors recent errors in UE logs"
    fi
}

# Check 9: UE Registration Status
check_ue_registration() {
    log "Checking UE registration status..."
    
    # Check if UE is registered by looking for registration success in logs
    if ssh $VM4 "grep -i 'registration successful' ue.log | tail -1" >/dev/null 2>&1; then
        record_status "UE_Registration" "PASS" "UE is registered"
    else
        record_status "UE_Registration" "FAIL" "UE is not registered"
    fi
}

# Check 10: PDU Session Status
check_pdu_session() {
    log "Checking PDU session status..."
    
    # Check if PDU session is established
    if ssh $VM4 "grep -i 'pdu session established' ue.log | tail -1" >/dev/null 2>&1; then
        record_status "PDU_Session" "PASS" "PDU session is established"
    else
        record_status "PDU_Session" "FAIL" "PDU session is not established"
    fi
}

# Main health check function
main() {
    print_header
    
    # Run all health checks
    check_vm_connectivity
    check_open5gs_services
    check_ueransim_processes
    check_network_ports
    check_sctp_associations
    check_database_connectivity
    check_system_resources
    check_recent_errors
    check_ue_registration
    check_pdu_session
    
    # Print results
    print_results
}

# Handle command line arguments
case "${1:-}" in
    "--continuous"|"-c")
        log "Starting continuous health monitoring..."
        while true; do
            main
            echo
            log "Waiting 60 seconds before next check..."
            sleep 60
        done
        ;;
    "--help"|"-h")
        echo "Usage: $0 [OPTIONS]"
        echo
        echo "Options:"
        echo "  -c, --continuous    Run health checks continuously"
        echo "  -h, --help          Show this help message"
        echo
        echo "Examples:"
        echo "  $0                  Run single health check"
        echo "  $0 -c               Run continuous health monitoring"
        ;;
    *)
        main
        ;;
esac 